#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@his-pc>
#
# Distributed under terms of the MIT license.


"""
经验之谈: 先完成一条测试路径, 测试完毕后, 再开发下一条.
比如这里: 先完成left-left 单旋转, 测试完毕后. 再开发left-right 双旋转.

参考 <<数据结构与算法分析-C语言版>> P82
示例也来自 <<数据结构与算法分析-C语言版>> P85, P86
"""

class AVLTreeNode:
    '''
        约定叶子节点高度为0, None节点高度为-1, 这样可以保证两者高度差为abs(0 - 1) = 1
    '''

    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right
        self.height = 0
        self.cnt = 1

    def insert(self, val):
        # 返回旋转后新的节点
        if val == self.val:
            self.cnt += 1
            return self

        if val < self.val:
            # 插入到左边
            # print('insert left')
            self.left = self._insert(self.left, val)
            return self.try_balance(val)
        else:
            # 插入到右边
            self.right = self._insert(self.right, val)
            return self.try_balance(val)

    def _insert(self, node, val):
        if node is None:
            return AVLTreeNode(val)
        else:
            return node.insert(val)

    def _height(self, node):
        if node is None:
            return -1

        return node.height

    def update_height(self):
        self.height = max(
            self._height(self.left),
            self._height(self.right)
        ) + 1

    def try_balance(self, val):
        if self._height(self.left) - self._height(self.right) > 1:
            if val < self.left.val:
                return self._rotate_left_left()
            else:
                return self._rotate_left_right()

        # rotate right
        if self._height(self.right) - self._height(self.left) > 1:
            if val > self.right.val:
                return self._rotate_right_right()
            else:
                return self._rotate_right_left()

        # 不需要旋转, 更新height即可
        self.update_height()
        return self

    def _rotate_left_left(self):
        # k2 > k1
        # k2 = self
        new_root = self.left        # self.left

        # print('rotate left-left k1, k2 [%s-%s, %s-%s]' % (k1.val, k1.height, k2.val, k2.height))

        self.left = new_root.right
        new_root.right = self

        # 注意这里先更新k2的高度, 再更新k1的高度(因为k1是k2的parent, parent的高度根据子节点高度更新)
        self.update_height()
        new_root.update_height()

        return new_root

    def _rotate_right_right(self):
        # k1 < k2, 旋转后k2将作为新的parent
        k1 = self
        k2 = k1.right       # self.right

        # print('rotate right-right k1, k2 [%s-%s, %s-%s]' % (k1.val, k1.height, k2.val, k2.height))

        k1.right = k2.left
        k2.left = k1

        # 重新设置 k1, k2 balance_factor
        k1.update_height()
        k2.update_height()

        return k2

    def _rotate_left_right(self):
        # 双旋 double rotate: k3 > k2 > k1
        # 这里可以分割成两个单旋来处理
        k3 = self
        k1 = self.left
        k2 = self.left.right

        # 验证关于k2左右子树高度的特征
        assert k2.height == 0 or (self._height(k2.left) != self._height(k2.right))

        # k3.left = k3.left._rotate_right_right()
        k3.left = k1._rotate_right_right()
        return k3._rotate_left_left()

    def _rotate_right_left(self):
        # k3 > k2 > k1
        # 这里还是手动处理, 为了以后查看.
        k1 = self
        k3 = self.right
        k2 = self.right.left

        # 后面内容和rotate_left_right 一样
        # print('rotate right-left (k3, k1, k2) (%s, %s, %s)' % (k3.val, k1.val, k2.val))

        # rotate
        k1.right = k2.left
        k3.left = k2.right

        k2.left = k1
        k2.right = k3

        k1.update_height()
        k3.update_height()
        k2.update_height()
        return k2

    def rec_show(self, cnt):
        if self.left:
            cnt = self.left.rec_show(cnt)

        print('%s:%s' % (self.val, self.height), end=' ', sep=' ')
        cnt += 1

        if self.right:
            cnt = self.right.rec_show(cnt)

        return cnt


class AVLTree:
    def __init__(self):
        self.root = None

    def insert(self, val):
        if self.root is None:
            self.root = AVLTreeNode(val)
            return self.root

        else:
            # self.root = self.root._insert(self.root, val)
            self.root  = self.root.insert(val)

    def show(self):
        if self.root:
            cnt = self.root.rec_show(cnt=0)
        print('\n=================== cnt %s =====================' % cnt)


def test_avl_insert():
    # =================== test left-left rotate =====================
    tree = AVLTree()

    for i in range(8, -1, -1):
        tree.insert(i)

    print('root %s' % tree.root.val)
    print('root.left %s' % tree.root.left.val)
    print('root.right %s' % tree.root.right.val)
    tree.show()

    # =================== test right-right rotate =====================
    tree = AVLTree()

    for i in range(8+1):
        tree.insert(i)

    print('root %s' % tree.root.val)
    print('root.left %s' % tree.root.left.val)
    print('root.right %s' % tree.root.right.val)
    tree.show()

    # =================== test left-right rotate =====================
    tree = AVLTree()
    tree.insert(10)
    tree.insert(5)
    tree.insert(7)

    print('root %s' % tree.root.val)
    print('root.left %s' % tree.root.left.val)
    print('root.right %s' % tree.root.right.val)
    tree.show()

    # =================== test right-left rotate =====================
    tree = AVLTree()
    tree.insert(5)
    tree.insert(10)
    tree.insert(7)

    print('root %s' % tree.root.val)
    print('root.left %s' % tree.root.left.val)
    print('root.right %s' % tree.root.right.val)
    tree.show()

    # # =================== test example rotate =====================
    tree = AVLTree()
    l = list(range(10000))
    import random
    random.shuffle(l)

    # print('=================== list %s ===================== ' % l)
    for i in l:
        tree.insert(i)
    tree.show()

    # =================== test example rotate =====================
    tree = AVLTree()
    tree.insert(4)
    tree.insert(2) ; tree.insert(6)
    tree.insert(1) ; tree.insert(3)
    tree.insert(5) ; tree.insert(7)
    tree.insert(16) ; tree.insert(15)
    tree.insert(14)
    tree.insert(13)
    tree.insert(12)
    tree.insert(11)
    tree.insert(10)
    tree.insert(8)
    tree.insert(9)

    print('root %s' % tree.root.val)
    print('root.left %s' % tree.root.left.val)
    print('root.right %s' % tree.root.right.val)
    tree.show()


def main():
    print("start main")
    test_avl_insert()

if __name__ == "__main__":
    main()
